import time

import cv2
import numpy as np
import onnxruntime as ort


MODEL_PATH = 'yolov8n.onnx'
CONF_THRESHOLD = 0.5
IOU_THRESHOLD = 0.1
CLASSES = [
 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 
 'traffic light', 'fire hydrant', 'street sign', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 
 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'hat', 'backpack', 'umbrella', 
 'shoe', 'eye glasses', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 
 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'plate', 
 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 
 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'mirror', 
 'dining table', 'window', 'desk', 'toilet', 'door', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 
 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'blender', 'book', 'clock', 'vase',
 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
]


def nms(boxes, scores):
    # Sort by score
    sorted_indices = np.argsort(scores)[::-1]

    keep_boxes = []
    while sorted_indices.size > 0:
        # Pick the last box
        box_id = sorted_indices[0]
        keep_boxes.append(box_id)

        # Compute IoU of the picked box with the rest
        ious = compute_iou(boxes[box_id, :], boxes[sorted_indices[1:], :])

        # Remove boxes with IoU over the threshold
        keep_indices = np.where(ious < IOU_THRESHOLD)[0]

        # print(keep_indices.shape, sorted_indices.shape)
        sorted_indices = sorted_indices[keep_indices + 1]

    return keep_boxes


def compute_iou(box, boxes):
    # Compute xmin, ymin, xmax, ymax for both boxes
    xmin = np.maximum(box[0], boxes[:, 0])
    ymin = np.maximum(box[1], boxes[:, 1])
    xmax = np.minimum(box[2], boxes[:, 2])
    ymax = np.minimum(box[3], boxes[:, 3])

    # Compute intersection area
    intersection_area = np.maximum(0, xmax - xmin) * np.maximum(0, ymax - ymin)

    # Compute union area
    box_area = (box[2] - box[0]) * (box[3] - box[1])
    boxes_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    union_area = box_area + boxes_area - intersection_area

    # Compute IoU
    iou = intersection_area / union_area

    return iou


def xywh2xyxy(xywh):
    xyxy = np.copy(xywh)
    xyxy[..., 0] = xywh[..., 0] - xywh[..., 2] / 2  # top left x
    xyxy[..., 1] = xywh[..., 1] - xywh[..., 3] / 2  # top left y
    xyxy[..., 2] = xywh[..., 0] + xywh[..., 2] / 2  # bottom right x
    xyxy[..., 3] = xywh[..., 1] + xywh[..., 3] / 2  # bottom right y
    return xyxy


def draw_predictions(image, boxes, scores, class_ids, indices, fps):
    cv2.putText(image, f'FPS: {fps}', (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.60, [225, 255, 255], thickness=2)
    for (bbox, score, label) in zip(xywh2xyxy(boxes[indices]), scores[indices], class_ids[indices]):
        bbox = bbox.round().astype(np.int32).tolist()
        cls_id = int(label)
        cls = CLASSES[cls_id]
        color = (0,255,0)
        cv2.rectangle(image, tuple(bbox[:2]), tuple(bbox[2:]), color, 2)
        cv2.putText(image,
                    f'{cls}: {int(score*100)}%', (bbox[0], bbox[1] - 2),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.60, [225, 255, 255],
                    thickness=1)


def preprocess(image, input_shape):
  """
  Funkcja przetwarzania wstępnego realizuje następujące kroki:
    1. zmienia rozmiar obrazu wejściowego, korzystając ze zmiennej _input_shape_,
    2. jeśli wykorzystywana jest biblioteka OpenCV to konwertuje z formatu BGR do RGB
    3. skaluje wartości w obrazie do zakresu <0; 1>,
    4. zmienia układ wymiarów z _HWC_ na _CHW_ (z formatu _channel last_ na _channel first_)
    5. dodaje dodatkowy wymiar do tensora, tak aby był w formacie _NHWC_ (_N_=1)
    6. ustawia format danych na _float32_
  """
  ##### Student code #####

  input_tensor = None

  ########################

  return input_tensor


def main():
    ##### Student code #####

    providers = ...

    ort_session = ...

    ########################
    print(f'Model loaded with {providers[0]}')

    model_inputs = ort_session.get_inputs()
    input_names = [model_inputs[i].name for i in range(len(model_inputs))]
    input_shape = model_inputs[0].shape

    model_output = ort_session.get_outputs()
    output_names = [model_output[i].name for i in range(len(model_output))]

    input_height, input_width = input_shape[2:]

    print('Model warm-up...')
    for _ in range(10):
        _ = ort_session.run(output_names, {input_names[0]: np.random.rand(1, 3, input_height, input_width).astype(np.float32)})[0]
    print('Model ready!')

    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print("Cannot open camera")
        exit()

    prev_frame_time = new_frame_time = 0

    while True:
        ret, frame = cap.read()
        if not ret:
            print("Can't receive frame. Exiting...")
            break

        new_frame_time = time.time()
        image_height, image_width = frame.shape[:2]

        input_tensor = preprocess(frame, input_shape=(input_width, input_height))
        outputs = ort_session.run(output_names, {input_names[0]: input_tensor})[0]

        predictions = np.squeeze(outputs).T

        # Filter out object confidence scores below threshold
        scores = np.max(predictions[:, 4:], axis=1)
        predictions = predictions[scores > CONF_THRESHOLD, :]
        scores = scores[scores > CONF_THRESHOLD]

        # Get the class with the highest confidence
        class_ids = np.argmax(predictions[:, 4:], axis=1)

        # Get bounding boxes for each object
        boxes = predictions[:, :4]

        # Rescale box
        input_shape = np.array([input_width, input_height, input_width, input_height])
        boxes = np.divide(boxes, input_shape, dtype=np.float32)
        boxes *= np.array([image_width, image_height, image_width, image_height])
        boxes = boxes.astype(np.int32)

        # Apply non-maxima suppression to suppress weak, overlapping bounding boxes
        indices = nms(boxes, scores)

        # Calculate frames per second rate
        fps = round(1 / (new_frame_time - prev_frame_time), 1)

        draw_predictions(frame, boxes, scores, class_ids, indices, fps)
        cv2.imshow('Predictions', frame)

        if cv2.waitKey(1) == ord('q'):
           break

        prev_frame_time = new_frame_time


if __name__ == '__main__':
    main()